import numpy as np
import matplotlib.pyplot as plt
from torch import nn, optim
from torch.autograd import Variable
import torch

# 生成非线性的数据(生成最小-2, 最大2的200个递增数据), 后面是加入维度(二维)
x_data = np.linspace(-2, 2, 200)[:, np.newaxis]
noise = np.random.normal(0, 0.2, x_data.shape)
y_data = np.square(x_data) + noise
plt.scatter(x_data, y_data)
plt.show()

# 把 numpy 转化成 tensor 数据
x_data = torch.FloatTensor(x_data)
y_data = torch.FloatTensor(y_data)
inputs = Variable(x_data)
target = Variable(y_data)


# 构建神经网络模型
# 一般把网络中具有可学习参数的层放在__init__()中
class LinearRegression(nn.Module):
    # 定义网络结果
    def __init__(self):
        # 初始化 nn.Module
        super(LinearRegression, self).__init__()
        # 定义网络结构, 加入隐藏层
        # 带有一层的神经网络
        self.fc1 = nn.Linear(1, 10)
        # 设置激活函数
        self.tanh = nn.Tanh()
        # 输出层
        self.fc2 = nn.Linear(10, 2)

    # 定义网络计算
    def forward(self, x):
        x = self.fc1(x)
        x = self.tanh(x)
        x = self.fc2(x)
        return x

# 定义模型
model = LinearRegression()

# 定义代价函数
mse_loss = nn.MSELoss()

# 定义优化器, lr表示学习率
optimizer = optim.SGD(model.parameters(), lr=0.3)

# 分别为参数的名字和参数的数值
for name, parameters in model.named_parameters():
    print('name:{}, param:{}'.format(name, parameters))

# 模型训练
for i in range(1000):
    out = model(inputs)
    # 计算loss
    loss = mse_loss(out, target)

    # 梯度清 0
    optimizer.zero_grad()

    # 计算梯度, 不清零会导致梯度累加
    loss.backward()

    # 修改权值
    optimizer.step()

    if i % 200 == 0:
        print(i, loss.item())

# 查看回归线的准确度
y_pred = model(inputs)
plt.scatter(x_data, y_data)
plt.plot(x_data, y_pred.data.numpy(), 'r-', lw=3)
plt.show()


